{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "scratchpad",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "accelerator": "TPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sayakpaul/Training-BatchNorm-and-Only-BatchNorm/blob/master/Varying_Batch_Sizes.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "lIYdn1woOS1n",
        "colab": {}
      },
      "source": [
        "# Authenticate yourself to use the TPUs\n",
        "import os\n",
        "\n",
        "IS_COLAB_BACKEND = 'COLAB_GPU' in os.environ  # this is always set on Colab, the value is 0 or 1 depending on GPU presence\n",
        "if IS_COLAB_BACKEND:\n",
        "  from google.colab import auth\n",
        "  # Authenticates the Colab machine and also the TPU using your\n",
        "  # credentials so that they can access your private GCS buckets.\n",
        "  auth.authenticate_user()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jT02YRxfsQIN",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "0c8e772f-1f8e-406c-f41d-c9f406ecabca"
      },
      "source": [
        "# TensorFlow Imports\n",
        "import tensorflow as tf\n",
        "print(tf.__version__)"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "2.2.0\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2LUjPO4zseS8",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 731
        },
        "outputId": "0032f7d9-08c2-4266-8fd3-ac62fa52c9ea"
      },
      "source": [
        "# Detect hardware\n",
        "try:\n",
        "  tpu = tf.distribute.cluster_resolver.TPUClusterResolver() # TPU detection\n",
        "except ValueError:\n",
        "  tpu = None\n",
        "  gpus = tf.config.experimental.list_logical_devices(\"GPU\")\n",
        "    \n",
        "# Select appropriate distribution strategy\n",
        "if tpu:\n",
        "  tf.config.experimental_connect_to_cluster(tpu)\n",
        "  tf.tpu.experimental.initialize_tpu_system(tpu)\n",
        "  strategy = tf.distribute.experimental.TPUStrategy(tpu) # Going back and forth between TPU and host is expensive. Better to run 128 batches on the TPU before reporting back.\n",
        "  print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])  \n",
        "elif len(gpus) > 1:\n",
        "  strategy = tf.distribute.MirroredStrategy([gpu.name for gpu in gpus])\n",
        "  print('Running on multiple GPUs ', [gpu.name for gpu in gpus])\n",
        "elif len(gpus) == 1:\n",
        "  strategy = tf.distribute.get_strategy() # default strategy that works on CPU and single GPU\n",
        "  print('Running on single GPU ', gpus[0].name)\n",
        "else:\n",
        "  strategy = tf.distribute.get_strategy() # default strategy that works on CPU and single GPU\n",
        "  print('Running on CPU')\n",
        "print(\"Number of accelerators: \", strategy.num_replicas_in_sync)"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Initializing the TPU system: grpc://10.58.143.242:8470\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Initializing the TPU system: grpc://10.58.143.242:8470\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Clearing out eager caches\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Clearing out eager caches\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Finished initializing TPU system.\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Finished initializing TPU system.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Found TPU system:\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Found TPU system:\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Num TPU Cores: 8\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Num TPU Cores: 8\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Num TPU Workers: 1\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Num TPU Workers: 1\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Num TPU Cores Per Worker: 8\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Num TPU Cores Per Worker: 8\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Running on TPU  ['10.58.143.242:8470']\n",
            "Number of accelerators:  8\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JvrSynThsgVg",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install wandb -q\n",
        "import wandb\n",
        "wandb.login()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2kHU76rnteAI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!wget https://raw.githubusercontent.com/GoogleCloudPlatform/keras-idiomatic-programmer/master/zoo/resnet/resnet_cifar10.py"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kPTxwomhtqw6",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Other imports\n",
        "from tensorflow.keras.layers import *\n",
        "from tensorflow.keras.models import *\n",
        "from wandb.keras import WandbCallback\n",
        "import matplotlib.pyplot as plt\n",
        "import tensorflow as tf\n",
        "import resnet_cifar10\n",
        "import numpy as np\n",
        "import time\n",
        "\n",
        "# Random seed fixation\n",
        "tf.random.set_seed(666)\n",
        "np.random.seed(666)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "399jYCS2ts93",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def get_training_model():\n",
        "    # ResNet20\n",
        "    n = 2\n",
        "    depth =  n * 9 + 2\n",
        "    n_blocks = ((depth - 2) // 9) - 1\n",
        "\n",
        "    # The input tensor\n",
        "    inputs = Input(shape=(32, 32, 3))\n",
        "\n",
        "    # The Stem Convolution Group\n",
        "    x = resnet_cifar10.stem(inputs)\n",
        "\n",
        "    # The learner\n",
        "    x = resnet_cifar10.learner(x, n_blocks)\n",
        "\n",
        "    # The Classifier for 10 classes\n",
        "    outputs = resnet_cifar10.classifier(x, 10)\n",
        "\n",
        "    # Instantiate the Model\n",
        "    model = Model(inputs, outputs)\n",
        "\n",
        "    for layer in model.layers:\n",
        "        if not isinstance(layer, tf.keras.layers.BatchNormalization):\n",
        "            if hasattr(layer, \"trainable\"):\n",
        "                layer.trainable=False\n",
        "\n",
        "    model.compile(loss=\"sparse_categorical_crossentropy\", optimizer=\"adam\", metrics=[\"accuracy\"])\n",
        "    \n",
        "    return model"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VCiAgWu5t8N_",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Load the training set of CIFAR10\n",
        "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3YkRBD2Dt_ER",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "BATCH_SIZE = 128 * strategy.num_replicas_in_sync\n",
        "\n",
        "def normalize(image, label):\n",
        "    image = tf.image.convert_image_dtype(image, tf.float32)\n",
        "    label = tf.cast(label, tf.int32)\n",
        "    return image, label\n",
        "\n",
        "def augment(image,label):\n",
        "    image = tf.image.resize_with_crop_or_pad(image, 40, 40) # Add 8 pixels of padding\n",
        "    image = tf.image.random_crop(image, size=[32, 32, 3]) # Random crop back to 32x32\n",
        "    image = tf.image.random_brightness(image, max_delta=0.5) # Random brightness\n",
        "    image = tf.clip_by_value(image, 0., 1.)\n",
        "\n",
        "    return image, label\n",
        "\n",
        "train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
        "train_ds = (\n",
        "    train_ds\n",
        "    .shuffle(1024)\n",
        "    .map(normalize, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "    .map(augment, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "    .batch(BATCH_SIZE)\n",
        "    .prefetch(tf.data.experimental.AUTOTUNE)\n",
        ")\n",
        "\n",
        "test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test))\n",
        "test_ds = (\n",
        "    test_ds\n",
        "    .map(normalize, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "    .batch(BATCH_SIZE)\n",
        "    .prefetch(tf.data.experimental.AUTOTUNE)\n",
        ")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ygJlCWGquV0y",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "outputId": "fb1a64ba-05c1-404d-a4ff-a40a985fa6ff"
      },
      "source": [
        "# Train model with with a batch size of 1024\n",
        "wandb.init(project=\"training-bn-only\", id=\"bn-bs-{}-crctd-3\".format(BATCH_SIZE))\n",
        "\n",
        "with strategy.scope():\n",
        "    model = get_training_model()\n",
        "\n",
        "start = time.time()\n",
        "h = model.fit(train_ds,\n",
        "         validation_data=test_ds,\n",
        "         epochs=75,\n",
        "         callbacks=[WandbCallback()])\n",
        "end = time.time()\n",
        "wandb.log({\"training_time\": end - start})\n",
        "wandb.log({\"batch_size\": BATCH_SIZE})\n",
        "print(\"Network takes {:.3f} seconds to train\".format(end - start))"
      ],
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "                Logging results to <a href=\"https://wandb.com\" target=\"_blank\">Weights & Biases</a> <a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">(Documentation)</a>.<br/>\n",
              "                Project page: <a href=\"https://app.wandb.ai/sayakpaul/training-bn-only\" target=\"_blank\">https://app.wandb.ai/sayakpaul/training-bn-only</a><br/>\n",
              "                Run page: <a href=\"https://app.wandb.ai/sayakpaul/training-bn-only/runs/bn-bs-1024-crctd-3\" target=\"_blank\">https://app.wandb.ai/sayakpaul/training-bn-only/runs/bn-bs-1024-crctd-3</a><br/>\n",
              "            "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "Epoch 1/75\n",
            "49/49 [==============================] - 10s 212ms/step - accuracy: 0.1173 - loss: 2.3456 - val_accuracy: 0.1671 - val_loss: 2.2664\n",
            "Epoch 2/75\n",
            "49/49 [==============================] - 3s 68ms/step - accuracy: 0.1490 - loss: 2.2658 - val_accuracy: 0.1579 - val_loss: 2.2433\n",
            "Epoch 3/75\n",
            "49/49 [==============================] - 2s 50ms/step - accuracy: 0.1826 - loss: 2.2107 - val_accuracy: 0.1472 - val_loss: 2.2507\n",
            "Epoch 4/75\n",
            "49/49 [==============================] - 3s 68ms/step - accuracy: 0.1995 - loss: 2.1722 - val_accuracy: 0.1544 - val_loss: 2.2387\n",
            "Epoch 5/75\n",
            "49/49 [==============================] - 3s 70ms/step - accuracy: 0.2093 - loss: 2.1424 - val_accuracy: 0.1667 - val_loss: 2.2020\n",
            "Epoch 6/75\n",
            "49/49 [==============================] - 3s 69ms/step - accuracy: 0.2167 - loss: 2.1143 - val_accuracy: 0.1771 - val_loss: 2.1509\n",
            "Epoch 7/75\n",
            "49/49 [==============================] - 3s 68ms/step - accuracy: 0.2242 - loss: 2.0915 - val_accuracy: 0.1976 - val_loss: 2.0958\n",
            "Epoch 8/75\n",
            "49/49 [==============================] - 3s 67ms/step - accuracy: 0.2303 - loss: 2.0712 - val_accuracy: 0.2197 - val_loss: 2.0541\n",
            "Epoch 9/75\n",
            "49/49 [==============================] - 4s 78ms/step - accuracy: 0.2390 - loss: 2.0484 - val_accuracy: 0.2415 - val_loss: 2.0114\n",
            "Epoch 10/75\n",
            "49/49 [==============================] - 3s 70ms/step - accuracy: 0.2454 - loss: 2.0310 - val_accuracy: 0.2561 - val_loss: 1.9822\n",
            "Epoch 11/75\n",
            "49/49 [==============================] - 3s 68ms/step - accuracy: 0.2492 - loss: 2.0151 - val_accuracy: 0.2656 - val_loss: 1.9547\n",
            "Epoch 12/75\n",
            "49/49 [==============================] - 3s 66ms/step - accuracy: 0.2553 - loss: 1.9990 - val_accuracy: 0.2775 - val_loss: 1.9343\n",
            "Epoch 13/75\n",
            "49/49 [==============================] - 3s 67ms/step - accuracy: 0.2607 - loss: 1.9812 - val_accuracy: 0.2870 - val_loss: 1.9137\n",
            "Epoch 14/75\n",
            "49/49 [==============================] - 3s 70ms/step - accuracy: 0.2689 - loss: 1.9657 - val_accuracy: 0.2933 - val_loss: 1.8945\n",
            "Epoch 15/75\n",
            "49/49 [==============================] - 3s 67ms/step - accuracy: 0.2673 - loss: 1.9533 - val_accuracy: 0.2965 - val_loss: 1.8792\n",
            "Epoch 16/75\n",
            "49/49 [==============================] - 3s 67ms/step - accuracy: 0.2709 - loss: 1.9411 - val_accuracy: 0.3027 - val_loss: 1.8659\n",
            "Epoch 17/75\n",
            "49/49 [==============================] - 3s 67ms/step - accuracy: 0.2787 - loss: 1.9262 - val_accuracy: 0.3103 - val_loss: 1.8519\n",
            "Epoch 18/75\n",
            "49/49 [==============================] - 4s 77ms/step - accuracy: 0.2807 - loss: 1.9176 - val_accuracy: 0.3148 - val_loss: 1.8360\n",
            "Epoch 19/75\n",
            "49/49 [==============================] - 3s 67ms/step - accuracy: 0.2841 - loss: 1.9081 - val_accuracy: 0.3186 - val_loss: 1.8254\n",
            "Epoch 20/75\n",
            "49/49 [==============================] - 3s 67ms/step - accuracy: 0.2919 - loss: 1.8931 - val_accuracy: 0.3202 - val_loss: 1.8140\n",
            "Epoch 21/75\n",
            "49/49 [==============================] - 3s 68ms/step - accuracy: 0.2926 - loss: 1.8869 - val_accuracy: 0.3242 - val_loss: 1.8059\n",
            "Epoch 22/75\n",
            "49/49 [==============================] - 3s 67ms/step - accuracy: 0.2964 - loss: 1.8774 - val_accuracy: 0.3295 - val_loss: 1.7973\n",
            "Epoch 23/75\n",
            "49/49 [==============================] - 3s 68ms/step - accuracy: 0.2980 - loss: 1.8699 - val_accuracy: 0.3357 - val_loss: 1.7878\n",
            "Epoch 24/75\n",
            "49/49 [==============================] - 3s 66ms/step - accuracy: 0.3041 - loss: 1.8589 - val_accuracy: 0.3382 - val_loss: 1.7799\n",
            "Epoch 25/75\n",
            "49/49 [==============================] - 3s 66ms/step - accuracy: 0.3076 - loss: 1.8524 - val_accuracy: 0.3418 - val_loss: 1.7725\n",
            "Epoch 26/75\n",
            "49/49 [==============================] - 3s 68ms/step - accuracy: 0.3075 - loss: 1.8455 - val_accuracy: 0.3452 - val_loss: 1.7646\n",
            "Epoch 27/75\n",
            "49/49 [==============================] - 4s 80ms/step - accuracy: 0.3112 - loss: 1.8356 - val_accuracy: 0.3457 - val_loss: 1.7583\n",
            "Epoch 28/75\n",
            "49/49 [==============================] - 3s 67ms/step - accuracy: 0.3123 - loss: 1.8310 - val_accuracy: 0.3516 - val_loss: 1.7494\n",
            "Epoch 29/75\n",
            "49/49 [==============================] - 3s 68ms/step - accuracy: 0.3161 - loss: 1.8230 - val_accuracy: 0.3542 - val_loss: 1.7446\n",
            "Epoch 30/75\n",
            "49/49 [==============================] - 3s 71ms/step - accuracy: 0.3165 - loss: 1.8209 - val_accuracy: 0.3542 - val_loss: 1.7370\n",
            "Epoch 31/75\n",
            "49/49 [==============================] - 3s 70ms/step - accuracy: 0.3227 - loss: 1.8124 - val_accuracy: 0.3555 - val_loss: 1.7337\n",
            "Epoch 32/75\n",
            "49/49 [==============================] - 3s 69ms/step - accuracy: 0.3210 - loss: 1.8065 - val_accuracy: 0.3565 - val_loss: 1.7282\n",
            "Epoch 33/75\n",
            "49/49 [==============================] - 3s 67ms/step - accuracy: 0.3225 - loss: 1.8023 - val_accuracy: 0.3611 - val_loss: 1.7204\n",
            "Epoch 34/75\n",
            "49/49 [==============================] - 3s 67ms/step - accuracy: 0.3249 - loss: 1.7998 - val_accuracy: 0.3623 - val_loss: 1.7151\n",
            "Epoch 35/75\n",
            "49/49 [==============================] - 3s 69ms/step - accuracy: 0.3274 - loss: 1.7898 - val_accuracy: 0.3628 - val_loss: 1.7086\n",
            "Epoch 36/75\n",
            "49/49 [==============================] - 3s 68ms/step - accuracy: 0.3320 - loss: 1.7852 - val_accuracy: 0.3650 - val_loss: 1.7044\n",
            "Epoch 37/75\n",
            "49/49 [==============================] - 3s 68ms/step - accuracy: 0.3355 - loss: 1.7819 - val_accuracy: 0.3649 - val_loss: 1.7009\n",
            "Epoch 38/75\n",
            "49/49 [==============================] - 3s 68ms/step - accuracy: 0.3378 - loss: 1.7782 - val_accuracy: 0.3689 - val_loss: 1.6953\n",
            "Epoch 39/75\n",
            "49/49 [==============================] - 3s 69ms/step - accuracy: 0.3397 - loss: 1.7737 - val_accuracy: 0.3713 - val_loss: 1.6885\n",
            "Epoch 40/75\n",
            "49/49 [==============================] - 3s 69ms/step - accuracy: 0.3398 - loss: 1.7679 - val_accuracy: 0.3737 - val_loss: 1.6852\n",
            "Epoch 41/75\n",
            "49/49 [==============================] - 3s 68ms/step - accuracy: 0.3404 - loss: 1.7638 - val_accuracy: 0.3751 - val_loss: 1.6815\n",
            "Epoch 42/75\n",
            "49/49 [==============================] - 3s 68ms/step - accuracy: 0.3451 - loss: 1.7568 - val_accuracy: 0.3759 - val_loss: 1.6736\n",
            "Epoch 43/75\n",
            "49/49 [==============================] - 3s 67ms/step - accuracy: 0.3446 - loss: 1.7570 - val_accuracy: 0.3801 - val_loss: 1.6707\n",
            "Epoch 44/75\n",
            "49/49 [==============================] - 3s 67ms/step - accuracy: 0.3471 - loss: 1.7483 - val_accuracy: 0.3809 - val_loss: 1.6661\n",
            "Epoch 45/75\n",
            "49/49 [==============================] - 3s 68ms/step - accuracy: 0.3514 - loss: 1.7453 - val_accuracy: 0.3843 - val_loss: 1.6617\n",
            "Epoch 46/75\n",
            "49/49 [==============================] - 3s 66ms/step - accuracy: 0.3538 - loss: 1.7410 - val_accuracy: 0.3859 - val_loss: 1.6567\n",
            "Epoch 47/75\n",
            "49/49 [==============================] - 3s 69ms/step - accuracy: 0.3515 - loss: 1.7395 - val_accuracy: 0.3901 - val_loss: 1.6511\n",
            "Epoch 48/75\n",
            "49/49 [==============================] - 3s 68ms/step - accuracy: 0.3522 - loss: 1.7376 - val_accuracy: 0.3882 - val_loss: 1.6482\n",
            "Epoch 49/75\n",
            "49/49 [==============================] - 3s 69ms/step - accuracy: 0.3579 - loss: 1.7313 - val_accuracy: 0.3911 - val_loss: 1.6455\n",
            "Epoch 50/75\n",
            "49/49 [==============================] - 3s 69ms/step - accuracy: 0.3608 - loss: 1.7271 - val_accuracy: 0.3944 - val_loss: 1.6397\n",
            "Epoch 51/75\n",
            "49/49 [==============================] - 3s 68ms/step - accuracy: 0.3595 - loss: 1.7225 - val_accuracy: 0.3965 - val_loss: 1.6363\n",
            "Epoch 52/75\n",
            "49/49 [==============================] - 3s 67ms/step - accuracy: 0.3600 - loss: 1.7218 - val_accuracy: 0.3986 - val_loss: 1.6317\n",
            "Epoch 53/75\n",
            "49/49 [==============================] - 3s 67ms/step - accuracy: 0.3621 - loss: 1.7198 - val_accuracy: 0.3996 - val_loss: 1.6266\n",
            "Epoch 54/75\n",
            "49/49 [==============================] - 3s 67ms/step - accuracy: 0.3634 - loss: 1.7168 - val_accuracy: 0.4020 - val_loss: 1.6241\n",
            "Epoch 55/75\n",
            "49/49 [==============================] - 4s 77ms/step - accuracy: 0.3680 - loss: 1.7092 - val_accuracy: 0.4038 - val_loss: 1.6209\n",
            "Epoch 56/75\n",
            "49/49 [==============================] - 3s 66ms/step - accuracy: 0.3641 - loss: 1.7107 - val_accuracy: 0.4055 - val_loss: 1.6186\n",
            "Epoch 57/75\n",
            "49/49 [==============================] - 3s 67ms/step - accuracy: 0.3704 - loss: 1.7039 - val_accuracy: 0.4074 - val_loss: 1.6151\n",
            "Epoch 58/75\n",
            "49/49 [==============================] - 3s 69ms/step - accuracy: 0.3722 - loss: 1.7023 - val_accuracy: 0.4120 - val_loss: 1.6097\n",
            "Epoch 59/75\n",
            "49/49 [==============================] - 3s 68ms/step - accuracy: 0.3693 - loss: 1.7018 - val_accuracy: 0.4113 - val_loss: 1.6083\n",
            "Epoch 60/75\n",
            "49/49 [==============================] - 3s 68ms/step - accuracy: 0.3740 - loss: 1.6970 - val_accuracy: 0.4120 - val_loss: 1.6046\n",
            "Epoch 61/75\n",
            "49/49 [==============================] - 3s 68ms/step - accuracy: 0.3747 - loss: 1.6963 - val_accuracy: 0.4122 - val_loss: 1.6004\n",
            "Epoch 62/75\n",
            "49/49 [==============================] - 3s 68ms/step - accuracy: 0.3752 - loss: 1.6931 - val_accuracy: 0.4131 - val_loss: 1.5992\n",
            "Epoch 63/75\n",
            "49/49 [==============================] - 3s 70ms/step - accuracy: 0.3775 - loss: 1.6866 - val_accuracy: 0.4143 - val_loss: 1.5974\n",
            "Epoch 64/75\n",
            "49/49 [==============================] - 2s 50ms/step - accuracy: 0.3788 - loss: 1.6828 - val_accuracy: 0.4144 - val_loss: 1.5979\n",
            "Epoch 65/75\n",
            "49/49 [==============================] - 3s 69ms/step - accuracy: 0.3761 - loss: 1.6861 - val_accuracy: 0.4182 - val_loss: 1.5909\n",
            "Epoch 66/75\n",
            "49/49 [==============================] - 4s 77ms/step - accuracy: 0.3828 - loss: 1.6780 - val_accuracy: 0.4185 - val_loss: 1.5874\n",
            "Epoch 67/75\n",
            "49/49 [==============================] - 2s 50ms/step - accuracy: 0.3829 - loss: 1.6785 - val_accuracy: 0.4184 - val_loss: 1.5884\n",
            "Epoch 68/75\n",
            "49/49 [==============================] - 3s 70ms/step - accuracy: 0.3835 - loss: 1.6745 - val_accuracy: 0.4195 - val_loss: 1.5831\n",
            "Epoch 69/75\n",
            "49/49 [==============================] - 3s 67ms/step - accuracy: 0.3834 - loss: 1.6758 - val_accuracy: 0.4219 - val_loss: 1.5796\n",
            "Epoch 70/75\n",
            "49/49 [==============================] - 3s 69ms/step - accuracy: 0.3839 - loss: 1.6700 - val_accuracy: 0.4226 - val_loss: 1.5763\n",
            "Epoch 71/75\n",
            "49/49 [==============================] - 3s 69ms/step - accuracy: 0.3825 - loss: 1.6694 - val_accuracy: 0.4249 - val_loss: 1.5734\n",
            "Epoch 72/75\n",
            "49/49 [==============================] - 3s 52ms/step - accuracy: 0.3853 - loss: 1.6680 - val_accuracy: 0.4249 - val_loss: 1.5743\n",
            "Epoch 73/75\n",
            "49/49 [==============================] - 2s 50ms/step - accuracy: 0.3886 - loss: 1.6619 - val_accuracy: 0.4241 - val_loss: 1.5759\n",
            "Epoch 74/75\n",
            "49/49 [==============================] - 2s 49ms/step - accuracy: 0.3862 - loss: 1.6649 - val_accuracy: 0.4239 - val_loss: 1.5763\n",
            "Epoch 75/75\n",
            "49/49 [==============================] - 3s 67ms/step - accuracy: 0.3889 - loss: 1.6625 - val_accuracy: 0.4255 - val_loss: 1.5718\n",
            "Network takes 308.746 seconds to train\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nEzu70uj0xcH",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "outputId": "3fbf7bb8-0899-4ca3-e94f-13ff534f34d7"
      },
      "source": [
        "BATCH_SIZE = 32 * strategy.num_replicas_in_sync\n",
        "\n",
        "train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
        "train_ds = (\n",
        "    train_ds\n",
        "    .shuffle(1024)\n",
        "    .map(normalize, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "    .map(augment, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "    .batch(BATCH_SIZE)\n",
        "    .prefetch(tf.data.experimental.AUTOTUNE)\n",
        ")\n",
        "\n",
        "test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test))\n",
        "test_ds = (\n",
        "    test_ds\n",
        "    .map(normalize, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "    .batch(BATCH_SIZE)\n",
        "    .prefetch(tf.data.experimental.AUTOTUNE)\n",
        ")\n",
        "\n",
        "# Train model with with a batch size of 256\n",
        "wandb.init(project=\"training-bn-only\", id=\"bn-bs-{}-crctd\".format(BATCH_SIZE))\n",
        "\n",
        "with strategy.scope():\n",
        "    model = get_training_model()\n",
        "\n",
        "start = time.time()\n",
        "h = model.fit(train_ds,\n",
        "         validation_data=test_ds,\n",
        "         epochs=75,\n",
        "         callbacks=[WandbCallback()])\n",
        "end = time.time()\n",
        "wandb.log({\"training_time\": end - start})\n",
        "wandb.log({\"batch_size\": BATCH_SIZE})\n",
        "print(\"Network takes {:.3f} seconds to train\".format(end - start))"
      ],
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "                Logging results to <a href=\"https://wandb.com\" target=\"_blank\">Weights & Biases</a> <a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">(Documentation)</a>.<br/>\n",
              "                Project page: <a href=\"https://app.wandb.ai/sayakpaul/training-bn-only\" target=\"_blank\">https://app.wandb.ai/sayakpaul/training-bn-only</a><br/>\n",
              "                Run page: <a href=\"https://app.wandb.ai/sayakpaul/training-bn-only/runs/bn-bs-256-crctd\" target=\"_blank\">https://app.wandb.ai/sayakpaul/training-bn-only/runs/bn-bs-256-crctd</a><br/>\n",
              "            "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "Epoch 1/75\n",
            "196/196 [==============================] - 14s 73ms/step - accuracy: 0.1413 - loss: 2.3104 - val_accuracy: 0.1195 - val_loss: 2.3373\n",
            "Epoch 2/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.1990 - loss: 2.1573 - val_accuracy: 0.2093 - val_loss: 2.0769\n",
            "Epoch 3/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.2258 - loss: 2.0839 - val_accuracy: 0.2582 - val_loss: 1.9708\n",
            "Epoch 4/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.2423 - loss: 2.0276 - val_accuracy: 0.2830 - val_loss: 1.9150\n",
            "Epoch 5/75\n",
            "196/196 [==============================] - 7s 37ms/step - accuracy: 0.2542 - loss: 1.9851 - val_accuracy: 0.2959 - val_loss: 1.8749\n",
            "Epoch 6/75\n",
            "196/196 [==============================] - 7s 37ms/step - accuracy: 0.2678 - loss: 1.9473 - val_accuracy: 0.3068 - val_loss: 1.8418\n",
            "Epoch 7/75\n",
            "196/196 [==============================] - 7s 37ms/step - accuracy: 0.2756 - loss: 1.9216 - val_accuracy: 0.3144 - val_loss: 1.8185\n",
            "Epoch 8/75\n",
            "196/196 [==============================] - 7s 37ms/step - accuracy: 0.2879 - loss: 1.8968 - val_accuracy: 0.3210 - val_loss: 1.7906\n",
            "Epoch 9/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.2960 - loss: 1.8717 - val_accuracy: 0.3270 - val_loss: 1.7687\n",
            "Epoch 10/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.3045 - loss: 1.8581 - val_accuracy: 0.3360 - val_loss: 1.7451\n",
            "Epoch 11/75\n",
            "196/196 [==============================] - 8s 41ms/step - accuracy: 0.3143 - loss: 1.8348 - val_accuracy: 0.3389 - val_loss: 1.7291\n",
            "Epoch 12/75\n",
            "196/196 [==============================] - 8s 39ms/step - accuracy: 0.3179 - loss: 1.8184 - val_accuracy: 0.3504 - val_loss: 1.7109\n",
            "Epoch 13/75\n",
            "196/196 [==============================] - 8s 39ms/step - accuracy: 0.3236 - loss: 1.8025 - val_accuracy: 0.3590 - val_loss: 1.6983\n",
            "Epoch 14/75\n",
            "196/196 [==============================] - 8s 40ms/step - accuracy: 0.3312 - loss: 1.7887 - val_accuracy: 0.3691 - val_loss: 1.6775\n",
            "Epoch 15/75\n",
            "196/196 [==============================] - 8s 38ms/step - accuracy: 0.3316 - loss: 1.7785 - val_accuracy: 0.3727 - val_loss: 1.6680\n",
            "Epoch 16/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.3402 - loss: 1.7627 - val_accuracy: 0.3811 - val_loss: 1.6548\n",
            "Epoch 17/75\n",
            "196/196 [==============================] - 8s 38ms/step - accuracy: 0.3454 - loss: 1.7544 - val_accuracy: 0.3877 - val_loss: 1.6394\n",
            "Epoch 18/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.3495 - loss: 1.7456 - val_accuracy: 0.3941 - val_loss: 1.6304\n",
            "Epoch 19/75\n",
            "196/196 [==============================] - 8s 38ms/step - accuracy: 0.3521 - loss: 1.7345 - val_accuracy: 0.3936 - val_loss: 1.6215\n",
            "Epoch 20/75\n",
            "196/196 [==============================] - 8s 39ms/step - accuracy: 0.3553 - loss: 1.7275 - val_accuracy: 0.4009 - val_loss: 1.6136\n",
            "Epoch 21/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.3578 - loss: 1.7210 - val_accuracy: 0.4014 - val_loss: 1.6080\n",
            "Epoch 22/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.3620 - loss: 1.7106 - val_accuracy: 0.4096 - val_loss: 1.5941\n",
            "Epoch 23/75\n",
            "196/196 [==============================] - 8s 39ms/step - accuracy: 0.3643 - loss: 1.7073 - val_accuracy: 0.4114 - val_loss: 1.5918\n",
            "Epoch 24/75\n",
            "196/196 [==============================] - 8s 41ms/step - accuracy: 0.3681 - loss: 1.7003 - val_accuracy: 0.4144 - val_loss: 1.5815\n",
            "Epoch 25/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.3721 - loss: 1.6949 - val_accuracy: 0.4126 - val_loss: 1.5779\n",
            "Epoch 26/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.3717 - loss: 1.6923 - val_accuracy: 0.4169 - val_loss: 1.5697\n",
            "Epoch 27/75\n",
            "196/196 [==============================] - 8s 39ms/step - accuracy: 0.3720 - loss: 1.6865 - val_accuracy: 0.4186 - val_loss: 1.5669\n",
            "Epoch 28/75\n",
            "196/196 [==============================] - 8s 41ms/step - accuracy: 0.3788 - loss: 1.6811 - val_accuracy: 0.4188 - val_loss: 1.5617\n",
            "Epoch 29/75\n",
            "196/196 [==============================] - 8s 40ms/step - accuracy: 0.3801 - loss: 1.6705 - val_accuracy: 0.4209 - val_loss: 1.5546\n",
            "Epoch 30/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.3830 - loss: 1.6685 - val_accuracy: 0.4259 - val_loss: 1.5476\n",
            "Epoch 31/75\n",
            "196/196 [==============================] - 8s 39ms/step - accuracy: 0.3832 - loss: 1.6646 - val_accuracy: 0.4300 - val_loss: 1.5442\n",
            "Epoch 32/75\n",
            "196/196 [==============================] - 8s 39ms/step - accuracy: 0.3875 - loss: 1.6587 - val_accuracy: 0.4311 - val_loss: 1.5375\n",
            "Epoch 33/75\n",
            "196/196 [==============================] - 8s 39ms/step - accuracy: 0.3900 - loss: 1.6547 - val_accuracy: 0.4353 - val_loss: 1.5357\n",
            "Epoch 34/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.3894 - loss: 1.6521 - val_accuracy: 0.4342 - val_loss: 1.5294\n",
            "Epoch 35/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.3911 - loss: 1.6488 - val_accuracy: 0.4373 - val_loss: 1.5253\n",
            "Epoch 36/75\n",
            "196/196 [==============================] - 8s 41ms/step - accuracy: 0.3949 - loss: 1.6428 - val_accuracy: 0.4404 - val_loss: 1.5244\n",
            "Epoch 37/75\n",
            "196/196 [==============================] - 8s 38ms/step - accuracy: 0.3944 - loss: 1.6408 - val_accuracy: 0.4425 - val_loss: 1.5186\n",
            "Epoch 38/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.3969 - loss: 1.6371 - val_accuracy: 0.4433 - val_loss: 1.5120\n",
            "Epoch 39/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.3972 - loss: 1.6327 - val_accuracy: 0.4489 - val_loss: 1.5053\n",
            "Epoch 40/75\n",
            "196/196 [==============================] - 7s 34ms/step - accuracy: 0.3999 - loss: 1.6294 - val_accuracy: 0.4416 - val_loss: 1.5056\n",
            "Epoch 41/75\n",
            "196/196 [==============================] - 8s 38ms/step - accuracy: 0.3983 - loss: 1.6267 - val_accuracy: 0.4463 - val_loss: 1.5004\n",
            "Epoch 42/75\n",
            "196/196 [==============================] - 8s 38ms/step - accuracy: 0.4027 - loss: 1.6209 - val_accuracy: 0.4435 - val_loss: 1.5003\n",
            "Epoch 43/75\n",
            "196/196 [==============================] - 8s 39ms/step - accuracy: 0.4048 - loss: 1.6192 - val_accuracy: 0.4499 - val_loss: 1.4962\n",
            "Epoch 44/75\n",
            "196/196 [==============================] - 8s 39ms/step - accuracy: 0.4084 - loss: 1.6133 - val_accuracy: 0.4523 - val_loss: 1.4893\n",
            "Epoch 45/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.4085 - loss: 1.6101 - val_accuracy: 0.4538 - val_loss: 1.4873\n",
            "Epoch 46/75\n",
            "196/196 [==============================] - 8s 39ms/step - accuracy: 0.4098 - loss: 1.6086 - val_accuracy: 0.4548 - val_loss: 1.4827\n",
            "Epoch 47/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.4113 - loss: 1.6060 - val_accuracy: 0.4526 - val_loss: 1.4825\n",
            "Epoch 48/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.4131 - loss: 1.6037 - val_accuracy: 0.4576 - val_loss: 1.4750\n",
            "Epoch 49/75\n",
            "196/196 [==============================] - 8s 39ms/step - accuracy: 0.4125 - loss: 1.6041 - val_accuracy: 0.4587 - val_loss: 1.4734\n",
            "Epoch 50/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.4125 - loss: 1.5993 - val_accuracy: 0.4622 - val_loss: 1.4657\n",
            "Epoch 51/75\n",
            "196/196 [==============================] - 7s 34ms/step - accuracy: 0.4182 - loss: 1.5907 - val_accuracy: 0.4605 - val_loss: 1.4677\n",
            "Epoch 52/75\n",
            "196/196 [==============================] - 8s 40ms/step - accuracy: 0.4210 - loss: 1.5881 - val_accuracy: 0.4649 - val_loss: 1.4623\n",
            "Epoch 53/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.4189 - loss: 1.5939 - val_accuracy: 0.4657 - val_loss: 1.4578\n",
            "Epoch 54/75\n",
            "196/196 [==============================] - 7s 34ms/step - accuracy: 0.4197 - loss: 1.5873 - val_accuracy: 0.4672 - val_loss: 1.4592\n",
            "Epoch 55/75\n",
            "196/196 [==============================] - 8s 39ms/step - accuracy: 0.4199 - loss: 1.5859 - val_accuracy: 0.4692 - val_loss: 1.4572\n",
            "Epoch 56/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.4197 - loss: 1.5805 - val_accuracy: 0.4702 - val_loss: 1.4530\n",
            "Epoch 57/75\n",
            "196/196 [==============================] - 8s 39ms/step - accuracy: 0.4195 - loss: 1.5813 - val_accuracy: 0.4704 - val_loss: 1.4496\n",
            "Epoch 58/75\n",
            "196/196 [==============================] - 7s 34ms/step - accuracy: 0.4242 - loss: 1.5791 - val_accuracy: 0.4683 - val_loss: 1.4572\n",
            "Epoch 59/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.4227 - loss: 1.5755 - val_accuracy: 0.4730 - val_loss: 1.4463\n",
            "Epoch 60/75\n",
            "196/196 [==============================] - 8s 38ms/step - accuracy: 0.4270 - loss: 1.5687 - val_accuracy: 0.4722 - val_loss: 1.4430\n",
            "Epoch 61/75\n",
            "196/196 [==============================] - 8s 41ms/step - accuracy: 0.4274 - loss: 1.5705 - val_accuracy: 0.4764 - val_loss: 1.4368\n",
            "Epoch 62/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.4259 - loss: 1.5713 - val_accuracy: 0.4766 - val_loss: 1.4348\n",
            "Epoch 63/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.4258 - loss: 1.5689 - val_accuracy: 0.4779 - val_loss: 1.4342\n",
            "Epoch 64/75\n",
            "196/196 [==============================] - 8s 39ms/step - accuracy: 0.4296 - loss: 1.5641 - val_accuracy: 0.4783 - val_loss: 1.4311\n",
            "Epoch 65/75\n",
            "196/196 [==============================] - 7s 34ms/step - accuracy: 0.4303 - loss: 1.5601 - val_accuracy: 0.4797 - val_loss: 1.4327\n",
            "Epoch 66/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.4318 - loss: 1.5540 - val_accuracy: 0.4828 - val_loss: 1.4238\n",
            "Epoch 67/75\n",
            "196/196 [==============================] - 7s 36ms/step - accuracy: 0.4327 - loss: 1.5571 - val_accuracy: 0.4815 - val_loss: 1.4246\n",
            "Epoch 68/75\n",
            "196/196 [==============================] - 8s 39ms/step - accuracy: 0.4331 - loss: 1.5533 - val_accuracy: 0.4774 - val_loss: 1.4235\n",
            "Epoch 69/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.4345 - loss: 1.5525 - val_accuracy: 0.4830 - val_loss: 1.4205\n",
            "Epoch 70/75\n",
            "196/196 [==============================] - 7s 34ms/step - accuracy: 0.4326 - loss: 1.5519 - val_accuracy: 0.4790 - val_loss: 1.4227\n",
            "Epoch 71/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.4354 - loss: 1.5463 - val_accuracy: 0.4809 - val_loss: 1.4204\n",
            "Epoch 72/75\n",
            "196/196 [==============================] - 7s 37ms/step - accuracy: 0.4350 - loss: 1.5468 - val_accuracy: 0.4870 - val_loss: 1.4122\n",
            "Epoch 73/75\n",
            "196/196 [==============================] - 7s 38ms/step - accuracy: 0.4344 - loss: 1.5486 - val_accuracy: 0.4861 - val_loss: 1.4097\n",
            "Epoch 74/75\n",
            "196/196 [==============================] - 8s 39ms/step - accuracy: 0.4358 - loss: 1.5440 - val_accuracy: 0.4881 - val_loss: 1.4131\n",
            "Epoch 75/75\n",
            "196/196 [==============================] - 8s 38ms/step - accuracy: 0.4403 - loss: 1.5420 - val_accuracy: 0.4894 - val_loss: 1.4064\n",
            "Network takes 616.974 seconds to train\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gghLkMpp2_dC",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "outputId": "bb872e36-64dc-4971-845b-e3c0bf539810"
      },
      "source": [
        "BATCH_SIZE = 64 * strategy.num_replicas_in_sync\n",
        "\n",
        "train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
        "train_ds = (\n",
        "    train_ds\n",
        "    .shuffle(1024)\n",
        "    .map(normalize, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "    .map(augment, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "    .batch(BATCH_SIZE)\n",
        "    .prefetch(tf.data.experimental.AUTOTUNE)\n",
        ")\n",
        "\n",
        "test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test))\n",
        "test_ds = (\n",
        "    test_ds\n",
        "    .map(normalize, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "    .batch(BATCH_SIZE)\n",
        "    .prefetch(tf.data.experimental.AUTOTUNE)\n",
        ")\n",
        "\n",
        "# Train model with with a batch size of 512\n",
        "wandb.init(project=\"training-bn-only\", id=\"bn-bs-{}\".format(BATCH_SIZE))\n",
        "\n",
        "with strategy.scope():\n",
        "    model = get_training_model()\n",
        "    \n",
        "start = time.time()\n",
        "h = model.fit(train_ds,\n",
        "         validation_data=test_ds,\n",
        "         epochs=75,\n",
        "         callbacks=[WandbCallback()])\n",
        "end = time.time()\n",
        "wandb.log({\"training_time\": end - start})\n",
        "wandb.log({\"batch_size\": BATCH_SIZE})\n",
        "print(\"Network takes {:.3f} seconds to train\".format(end - start))"
      ],
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "                Logging results to <a href=\"https://wandb.com\" target=\"_blank\">Weights & Biases</a> <a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">(Documentation)</a>.<br/>\n",
              "                Project page: <a href=\"https://app.wandb.ai/sayakpaul/training-bn-only\" target=\"_blank\">https://app.wandb.ai/sayakpaul/training-bn-only</a><br/>\n",
              "                Run page: <a href=\"https://app.wandb.ai/sayakpaul/training-bn-only/runs/bn-bs-512\" target=\"_blank\">https://app.wandb.ai/sayakpaul/training-bn-only/runs/bn-bs-512</a><br/>\n",
              "            "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "Epoch 1/75\n",
            "98/98 [==============================] - 13s 128ms/step - accuracy: 0.1247 - loss: 2.3250 - val_accuracy: 0.1369 - val_loss: 2.2511\n",
            "Epoch 2/75\n",
            "98/98 [==============================] - 5s 48ms/step - accuracy: 0.1804 - loss: 2.2067 - val_accuracy: 0.1500 - val_loss: 2.2039\n",
            "Epoch 3/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.2139 - loss: 2.1329 - val_accuracy: 0.2000 - val_loss: 2.1233\n",
            "Epoch 4/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.2286 - loss: 2.0866 - val_accuracy: 0.2191 - val_loss: 2.0634\n",
            "Epoch 5/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.2419 - loss: 2.0549 - val_accuracy: 0.2451 - val_loss: 2.0257\n",
            "Epoch 6/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.2489 - loss: 2.0307 - val_accuracy: 0.2680 - val_loss: 1.9882\n",
            "Epoch 7/75\n",
            "98/98 [==============================] - 5s 50ms/step - accuracy: 0.2560 - loss: 2.0125 - val_accuracy: 0.2801 - val_loss: 1.9637\n",
            "Epoch 8/75\n",
            "98/98 [==============================] - 5s 51ms/step - accuracy: 0.2620 - loss: 1.9944 - val_accuracy: 0.2900 - val_loss: 1.9401\n",
            "Epoch 9/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.2700 - loss: 1.9798 - val_accuracy: 0.3003 - val_loss: 1.9196\n",
            "Epoch 10/75\n",
            "98/98 [==============================] - 5s 48ms/step - accuracy: 0.2752 - loss: 1.9619 - val_accuracy: 0.3129 - val_loss: 1.8972\n",
            "Epoch 11/75\n",
            "98/98 [==============================] - 5s 48ms/step - accuracy: 0.2809 - loss: 1.9501 - val_accuracy: 0.3178 - val_loss: 1.8814\n",
            "Epoch 12/75\n",
            "98/98 [==============================] - 5s 48ms/step - accuracy: 0.2826 - loss: 1.9357 - val_accuracy: 0.3223 - val_loss: 1.8661\n",
            "Epoch 13/75\n",
            "98/98 [==============================] - 5s 48ms/step - accuracy: 0.2876 - loss: 1.9270 - val_accuracy: 0.3255 - val_loss: 1.8516\n",
            "Epoch 14/75\n",
            "98/98 [==============================] - 5s 55ms/step - accuracy: 0.2940 - loss: 1.9117 - val_accuracy: 0.3337 - val_loss: 1.8341\n",
            "Epoch 15/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.2970 - loss: 1.8988 - val_accuracy: 0.3362 - val_loss: 1.8228\n",
            "Epoch 16/75\n",
            "98/98 [==============================] - 5s 50ms/step - accuracy: 0.3035 - loss: 1.8849 - val_accuracy: 0.3402 - val_loss: 1.8100\n",
            "Epoch 17/75\n",
            "98/98 [==============================] - 5s 52ms/step - accuracy: 0.3071 - loss: 1.8715 - val_accuracy: 0.3493 - val_loss: 1.7884\n",
            "Epoch 18/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.3111 - loss: 1.8609 - val_accuracy: 0.3534 - val_loss: 1.7745\n",
            "Epoch 19/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.3135 - loss: 1.8495 - val_accuracy: 0.3555 - val_loss: 1.7646\n",
            "Epoch 20/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.3197 - loss: 1.8347 - val_accuracy: 0.3581 - val_loss: 1.7531\n",
            "Epoch 21/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.3239 - loss: 1.8236 - val_accuracy: 0.3591 - val_loss: 1.7416\n",
            "Epoch 22/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.3257 - loss: 1.8145 - val_accuracy: 0.3669 - val_loss: 1.7279\n",
            "Epoch 23/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.3318 - loss: 1.8053 - val_accuracy: 0.3693 - val_loss: 1.7186\n",
            "Epoch 24/75\n",
            "98/98 [==============================] - 5s 48ms/step - accuracy: 0.3358 - loss: 1.7924 - val_accuracy: 0.3733 - val_loss: 1.7073\n",
            "Epoch 25/75\n",
            "98/98 [==============================] - 5s 48ms/step - accuracy: 0.3377 - loss: 1.7871 - val_accuracy: 0.3788 - val_loss: 1.6969\n",
            "Epoch 26/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.3427 - loss: 1.7769 - val_accuracy: 0.3833 - val_loss: 1.6884\n",
            "Epoch 27/75\n",
            "98/98 [==============================] - 6s 56ms/step - accuracy: 0.3436 - loss: 1.7684 - val_accuracy: 0.3826 - val_loss: 1.6805\n",
            "Epoch 28/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.3494 - loss: 1.7627 - val_accuracy: 0.3832 - val_loss: 1.6732\n",
            "Epoch 29/75\n",
            "98/98 [==============================] - 5s 51ms/step - accuracy: 0.3510 - loss: 1.7536 - val_accuracy: 0.3919 - val_loss: 1.6628\n",
            "Epoch 30/75\n",
            "98/98 [==============================] - 5s 50ms/step - accuracy: 0.3550 - loss: 1.7459 - val_accuracy: 0.3935 - val_loss: 1.6553\n",
            "Epoch 31/75\n",
            "98/98 [==============================] - 5s 50ms/step - accuracy: 0.3553 - loss: 1.7427 - val_accuracy: 0.3944 - val_loss: 1.6486\n",
            "Epoch 32/75\n",
            "98/98 [==============================] - 5s 50ms/step - accuracy: 0.3597 - loss: 1.7365 - val_accuracy: 0.4004 - val_loss: 1.6404\n",
            "Epoch 33/75\n",
            "98/98 [==============================] - 5s 48ms/step - accuracy: 0.3634 - loss: 1.7266 - val_accuracy: 0.4019 - val_loss: 1.6375\n",
            "Epoch 34/75\n",
            "98/98 [==============================] - 5s 48ms/step - accuracy: 0.3639 - loss: 1.7231 - val_accuracy: 0.4055 - val_loss: 1.6267\n",
            "Epoch 35/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.3661 - loss: 1.7174 - val_accuracy: 0.4016 - val_loss: 1.6262\n",
            "Epoch 36/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.3673 - loss: 1.7140 - val_accuracy: 0.4022 - val_loss: 1.6248\n",
            "Epoch 37/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.3691 - loss: 1.7090 - val_accuracy: 0.4050 - val_loss: 1.6164\n",
            "Epoch 38/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.3710 - loss: 1.7043 - val_accuracy: 0.4066 - val_loss: 1.6109\n",
            "Epoch 39/75\n",
            "98/98 [==============================] - 4s 41ms/step - accuracy: 0.3723 - loss: 1.7009 - val_accuracy: 0.4055 - val_loss: 1.6152\n",
            "Epoch 40/75\n",
            "98/98 [==============================] - 5s 51ms/step - accuracy: 0.3737 - loss: 1.6978 - val_accuracy: 0.4126 - val_loss: 1.6010\n",
            "Epoch 41/75\n",
            "98/98 [==============================] - 5s 51ms/step - accuracy: 0.3739 - loss: 1.6935 - val_accuracy: 0.4132 - val_loss: 1.5979\n",
            "Epoch 42/75\n",
            "98/98 [==============================] - 6s 57ms/step - accuracy: 0.3798 - loss: 1.6878 - val_accuracy: 0.4130 - val_loss: 1.5935\n",
            "Epoch 43/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.3755 - loss: 1.6856 - val_accuracy: 0.4158 - val_loss: 1.5929\n",
            "Epoch 44/75\n",
            "98/98 [==============================] - 5s 51ms/step - accuracy: 0.3837 - loss: 1.6799 - val_accuracy: 0.4149 - val_loss: 1.5907\n",
            "Epoch 45/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.3838 - loss: 1.6762 - val_accuracy: 0.4156 - val_loss: 1.5845\n",
            "Epoch 46/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.3843 - loss: 1.6748 - val_accuracy: 0.4198 - val_loss: 1.5789\n",
            "Epoch 47/75\n",
            "98/98 [==============================] - 4s 40ms/step - accuracy: 0.3837 - loss: 1.6710 - val_accuracy: 0.4174 - val_loss: 1.5819\n",
            "Epoch 48/75\n",
            "98/98 [==============================] - 4s 39ms/step - accuracy: 0.3850 - loss: 1.6650 - val_accuracy: 0.4180 - val_loss: 1.5812\n",
            "Epoch 49/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.3862 - loss: 1.6656 - val_accuracy: 0.4195 - val_loss: 1.5778\n",
            "Epoch 50/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.3880 - loss: 1.6626 - val_accuracy: 0.4230 - val_loss: 1.5671\n",
            "Epoch 51/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.3904 - loss: 1.6597 - val_accuracy: 0.4223 - val_loss: 1.5660\n",
            "Epoch 52/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.3922 - loss: 1.6561 - val_accuracy: 0.4221 - val_loss: 1.5619\n",
            "Epoch 53/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.3912 - loss: 1.6553 - val_accuracy: 0.4266 - val_loss: 1.5537\n",
            "Epoch 54/75\n",
            "98/98 [==============================] - 4s 41ms/step - accuracy: 0.3938 - loss: 1.6468 - val_accuracy: 0.4273 - val_loss: 1.5580\n",
            "Epoch 55/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.3951 - loss: 1.6452 - val_accuracy: 0.4301 - val_loss: 1.5472\n",
            "Epoch 56/75\n",
            "98/98 [==============================] - 4s 40ms/step - accuracy: 0.3951 - loss: 1.6414 - val_accuracy: 0.4275 - val_loss: 1.5505\n",
            "Epoch 57/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.3981 - loss: 1.6407 - val_accuracy: 0.4338 - val_loss: 1.5467\n",
            "Epoch 58/75\n",
            "98/98 [==============================] - 6s 56ms/step - accuracy: 0.3987 - loss: 1.6419 - val_accuracy: 0.4336 - val_loss: 1.5384\n",
            "Epoch 59/75\n",
            "98/98 [==============================] - 4s 41ms/step - accuracy: 0.3995 - loss: 1.6358 - val_accuracy: 0.4328 - val_loss: 1.5417\n",
            "Epoch 60/75\n",
            "98/98 [==============================] - 4s 41ms/step - accuracy: 0.3991 - loss: 1.6362 - val_accuracy: 0.4311 - val_loss: 1.5464\n",
            "Epoch 61/75\n",
            "98/98 [==============================] - 4s 40ms/step - accuracy: 0.4008 - loss: 1.6321 - val_accuracy: 0.4312 - val_loss: 1.5469\n",
            "Epoch 62/75\n",
            "98/98 [==============================] - 4s 40ms/step - accuracy: 0.4026 - loss: 1.6276 - val_accuracy: 0.4313 - val_loss: 1.5573\n",
            "Epoch 63/75\n",
            "98/98 [==============================] - 4s 40ms/step - accuracy: 0.4040 - loss: 1.6244 - val_accuracy: 0.4356 - val_loss: 1.5406\n",
            "Epoch 64/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.4056 - loss: 1.6258 - val_accuracy: 0.4418 - val_loss: 1.5309\n",
            "Epoch 65/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.4022 - loss: 1.6251 - val_accuracy: 0.4417 - val_loss: 1.5276\n",
            "Epoch 66/75\n",
            "98/98 [==============================] - 4s 41ms/step - accuracy: 0.4047 - loss: 1.6221 - val_accuracy: 0.4400 - val_loss: 1.5279\n",
            "Epoch 67/75\n",
            "98/98 [==============================] - 4s 40ms/step - accuracy: 0.4092 - loss: 1.6157 - val_accuracy: 0.4384 - val_loss: 1.5302\n",
            "Epoch 68/75\n",
            "98/98 [==============================] - 4s 40ms/step - accuracy: 0.4074 - loss: 1.6156 - val_accuracy: 0.4427 - val_loss: 1.5314\n",
            "Epoch 69/75\n",
            "98/98 [==============================] - 5s 50ms/step - accuracy: 0.4081 - loss: 1.6131 - val_accuracy: 0.4431 - val_loss: 1.5239\n",
            "Epoch 70/75\n",
            "98/98 [==============================] - 5s 50ms/step - accuracy: 0.4081 - loss: 1.6096 - val_accuracy: 0.4427 - val_loss: 1.5218\n",
            "Epoch 71/75\n",
            "98/98 [==============================] - 5s 49ms/step - accuracy: 0.4134 - loss: 1.6091 - val_accuracy: 0.4459 - val_loss: 1.5168\n",
            "Epoch 72/75\n",
            "98/98 [==============================] - 5s 48ms/step - accuracy: 0.4117 - loss: 1.6078 - val_accuracy: 0.4470 - val_loss: 1.5103\n",
            "Epoch 73/75\n",
            "98/98 [==============================] - 6s 56ms/step - accuracy: 0.4137 - loss: 1.6006 - val_accuracy: 0.4444 - val_loss: 1.5099\n",
            "Epoch 74/75\n",
            "98/98 [==============================] - 5s 50ms/step - accuracy: 0.4141 - loss: 1.6028 - val_accuracy: 0.4454 - val_loss: 1.5076\n",
            "Epoch 75/75\n",
            "98/98 [==============================] - 4s 41ms/step - accuracy: 0.4138 - loss: 1.6040 - val_accuracy: 0.4448 - val_loss: 1.5227\n",
            "Network takes 409.876 seconds to train\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ju1yCcV03u4w",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "outputId": "34af237d-1044-4149-8ba3-63d109d2d06d"
      },
      "source": [
        "BATCH_SIZE = 256 * strategy.num_replicas_in_sync\n",
        "\n",
        "train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
        "train_ds = (\n",
        "    train_ds\n",
        "    .shuffle(1024)\n",
        "    .map(normalize, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "    .map(augment, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "    .batch(BATCH_SIZE)\n",
        "    .prefetch(tf.data.experimental.AUTOTUNE)\n",
        ")\n",
        "\n",
        "test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test))\n",
        "test_ds = (\n",
        "    test_ds\n",
        "    .map(normalize, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "    .batch(BATCH_SIZE)\n",
        "    .prefetch(tf.data.experimental.AUTOTUNE)\n",
        ")\n",
        "\n",
        "# Train model with with a batch size of 2048\n",
        "wandb.init(project=\"training-bn-only\", id=\"bn-bs-{}\".format(BATCH_SIZE))\n",
        "\n",
        "with strategy.scope():\n",
        "    model = get_training_model()\n",
        "    \n",
        "start = time.time()\n",
        "h = model.fit(train_ds,\n",
        "         validation_data=test_ds,\n",
        "         epochs=75,\n",
        "         callbacks=[WandbCallback()])\n",
        "end = time.time()\n",
        "wandb.log({\"training_time\": end - start})\n",
        "wandb.log({\"batch_size\": BATCH_SIZE})\n",
        "print(\"Network takes {:.3f} seconds to train\".format(end - start))"
      ],
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "                Logging results to <a href=\"https://wandb.com\" target=\"_blank\">Weights & Biases</a> <a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">(Documentation)</a>.<br/>\n",
              "                Project page: <a href=\"https://app.wandb.ai/sayakpaul/training-bn-only\" target=\"_blank\">https://app.wandb.ai/sayakpaul/training-bn-only</a><br/>\n",
              "                Run page: <a href=\"https://app.wandb.ai/sayakpaul/training-bn-only/runs/bn-bs-2048\" target=\"_blank\">https://app.wandb.ai/sayakpaul/training-bn-only/runs/bn-bs-2048</a><br/>\n",
              "            "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "Epoch 1/75\n",
            "25/25 [==============================] - 10s 419ms/step - accuracy: 0.1122 - loss: 2.3638 - val_accuracy: 0.1204 - val_loss: 2.3491\n",
            "Epoch 2/75\n",
            "25/25 [==============================] - 3s 124ms/step - accuracy: 0.1297 - loss: 2.3071 - val_accuracy: 0.1525 - val_loss: 2.2702\n",
            "Epoch 3/75\n",
            "25/25 [==============================] - 2s 90ms/step - accuracy: 0.1468 - loss: 2.2638 - val_accuracy: 0.1212 - val_loss: 2.2789\n",
            "Epoch 4/75\n",
            "25/25 [==============================] - 2s 88ms/step - accuracy: 0.1635 - loss: 2.2289 - val_accuracy: 0.1125 - val_loss: 2.2900\n",
            "Epoch 5/75\n",
            "25/25 [==============================] - 2s 90ms/step - accuracy: 0.1734 - loss: 2.1988 - val_accuracy: 0.1114 - val_loss: 2.3024\n",
            "Epoch 6/75\n",
            "25/25 [==============================] - 2s 87ms/step - accuracy: 0.1888 - loss: 2.1718 - val_accuracy: 0.1168 - val_loss: 2.3023\n",
            "Epoch 7/75\n",
            "25/25 [==============================] - 2s 91ms/step - accuracy: 0.2003 - loss: 2.1492 - val_accuracy: 0.1270 - val_loss: 2.2892\n",
            "Epoch 8/75\n",
            "25/25 [==============================] - 3s 123ms/step - accuracy: 0.2052 - loss: 2.1271 - val_accuracy: 0.1406 - val_loss: 2.2622\n",
            "Epoch 9/75\n",
            "25/25 [==============================] - 3s 124ms/step - accuracy: 0.2121 - loss: 2.1103 - val_accuracy: 0.1570 - val_loss: 2.2252\n",
            "Epoch 10/75\n",
            "25/25 [==============================] - 3s 125ms/step - accuracy: 0.2178 - loss: 2.0924 - val_accuracy: 0.1702 - val_loss: 2.1825\n",
            "Epoch 11/75\n",
            "25/25 [==============================] - 3s 125ms/step - accuracy: 0.2229 - loss: 2.0779 - val_accuracy: 0.1827 - val_loss: 2.1432\n",
            "Epoch 12/75\n",
            "25/25 [==============================] - 3s 120ms/step - accuracy: 0.2288 - loss: 2.0637 - val_accuracy: 0.1941 - val_loss: 2.1056\n",
            "Epoch 13/75\n",
            "25/25 [==============================] - 3s 123ms/step - accuracy: 0.2353 - loss: 2.0501 - val_accuracy: 0.2086 - val_loss: 2.0796\n",
            "Epoch 14/75\n",
            "25/25 [==============================] - 3s 123ms/step - accuracy: 0.2373 - loss: 2.0399 - val_accuracy: 0.2189 - val_loss: 2.0458\n",
            "Epoch 15/75\n",
            "25/25 [==============================] - 3s 137ms/step - accuracy: 0.2448 - loss: 2.0256 - val_accuracy: 0.2287 - val_loss: 2.0209\n",
            "Epoch 16/75\n",
            "25/25 [==============================] - 3s 128ms/step - accuracy: 0.2439 - loss: 2.0180 - val_accuracy: 0.2407 - val_loss: 2.0012\n",
            "Epoch 17/75\n",
            "25/25 [==============================] - 3s 123ms/step - accuracy: 0.2507 - loss: 2.0067 - val_accuracy: 0.2455 - val_loss: 1.9827\n",
            "Epoch 18/75\n",
            "25/25 [==============================] - 3s 125ms/step - accuracy: 0.2561 - loss: 1.9944 - val_accuracy: 0.2566 - val_loss: 1.9667\n",
            "Epoch 19/75\n",
            "25/25 [==============================] - 4s 162ms/step - accuracy: 0.2575 - loss: 1.9873 - val_accuracy: 0.2617 - val_loss: 1.9549\n",
            "Epoch 20/75\n",
            "25/25 [==============================] - 3s 122ms/step - accuracy: 0.2608 - loss: 1.9789 - val_accuracy: 0.2706 - val_loss: 1.9453\n",
            "Epoch 21/75\n",
            "25/25 [==============================] - 3s 124ms/step - accuracy: 0.2627 - loss: 1.9729 - val_accuracy: 0.2749 - val_loss: 1.9316\n",
            "Epoch 22/75\n",
            "25/25 [==============================] - 3s 124ms/step - accuracy: 0.2662 - loss: 1.9634 - val_accuracy: 0.2793 - val_loss: 1.9186\n",
            "Epoch 23/75\n",
            "25/25 [==============================] - 3s 122ms/step - accuracy: 0.2720 - loss: 1.9531 - val_accuracy: 0.2857 - val_loss: 1.9064\n",
            "Epoch 24/75\n",
            "25/25 [==============================] - 3s 122ms/step - accuracy: 0.2750 - loss: 1.9439 - val_accuracy: 0.2914 - val_loss: 1.9023\n",
            "Epoch 25/75\n",
            "25/25 [==============================] - 3s 123ms/step - accuracy: 0.2746 - loss: 1.9394 - val_accuracy: 0.2934 - val_loss: 1.8931\n",
            "Epoch 26/75\n",
            "25/25 [==============================] - 3s 124ms/step - accuracy: 0.2805 - loss: 1.9325 - val_accuracy: 0.2980 - val_loss: 1.8833\n",
            "Epoch 27/75\n",
            "25/25 [==============================] - 3s 122ms/step - accuracy: 0.2820 - loss: 1.9258 - val_accuracy: 0.3035 - val_loss: 1.8746\n",
            "Epoch 28/75\n",
            "25/25 [==============================] - 3s 133ms/step - accuracy: 0.2830 - loss: 1.9176 - val_accuracy: 0.3018 - val_loss: 1.8681\n",
            "Epoch 29/75\n",
            "25/25 [==============================] - 3s 133ms/step - accuracy: 0.2897 - loss: 1.9086 - val_accuracy: 0.3042 - val_loss: 1.8630\n",
            "Epoch 30/75\n",
            "25/25 [==============================] - 3s 128ms/step - accuracy: 0.2928 - loss: 1.9019 - val_accuracy: 0.3080 - val_loss: 1.8584\n",
            "Epoch 31/75\n",
            "25/25 [==============================] - 3s 124ms/step - accuracy: 0.2936 - loss: 1.8992 - val_accuracy: 0.3112 - val_loss: 1.8506\n",
            "Epoch 32/75\n",
            "25/25 [==============================] - 3s 126ms/step - accuracy: 0.2923 - loss: 1.8931 - val_accuracy: 0.3151 - val_loss: 1.8428\n",
            "Epoch 33/75\n",
            "25/25 [==============================] - 3s 129ms/step - accuracy: 0.2979 - loss: 1.8846 - val_accuracy: 0.3168 - val_loss: 1.8359\n",
            "Epoch 34/75\n",
            "25/25 [==============================] - 3s 123ms/step - accuracy: 0.3018 - loss: 1.8784 - val_accuracy: 0.3184 - val_loss: 1.8326\n",
            "Epoch 35/75\n",
            "25/25 [==============================] - 3s 122ms/step - accuracy: 0.3024 - loss: 1.8726 - val_accuracy: 0.3198 - val_loss: 1.8237\n",
            "Epoch 36/75\n",
            "25/25 [==============================] - 4s 159ms/step - accuracy: 0.3019 - loss: 1.8688 - val_accuracy: 0.3216 - val_loss: 1.8217\n",
            "Epoch 37/75\n",
            "25/25 [==============================] - 3s 124ms/step - accuracy: 0.3036 - loss: 1.8631 - val_accuracy: 0.3220 - val_loss: 1.8141\n",
            "Epoch 38/75\n",
            "25/25 [==============================] - 3s 126ms/step - accuracy: 0.3081 - loss: 1.8571 - val_accuracy: 0.3248 - val_loss: 1.8092\n",
            "Epoch 39/75\n",
            "25/25 [==============================] - 3s 121ms/step - accuracy: 0.3090 - loss: 1.8534 - val_accuracy: 0.3266 - val_loss: 1.8042\n",
            "Epoch 40/75\n",
            "25/25 [==============================] - 3s 122ms/step - accuracy: 0.3111 - loss: 1.8474 - val_accuracy: 0.3257 - val_loss: 1.8024\n",
            "Epoch 41/75\n",
            "25/25 [==============================] - 3s 124ms/step - accuracy: 0.3114 - loss: 1.8432 - val_accuracy: 0.3276 - val_loss: 1.7979\n",
            "Epoch 42/75\n",
            "25/25 [==============================] - 3s 123ms/step - accuracy: 0.3140 - loss: 1.8353 - val_accuracy: 0.3312 - val_loss: 1.7943\n",
            "Epoch 43/75\n",
            "25/25 [==============================] - 3s 122ms/step - accuracy: 0.3153 - loss: 1.8319 - val_accuracy: 0.3355 - val_loss: 1.7851\n",
            "Epoch 44/75\n",
            "25/25 [==============================] - 3s 121ms/step - accuracy: 0.3194 - loss: 1.8252 - val_accuracy: 0.3358 - val_loss: 1.7774\n",
            "Epoch 45/75\n",
            "25/25 [==============================] - 3s 125ms/step - accuracy: 0.3178 - loss: 1.8239 - val_accuracy: 0.3364 - val_loss: 1.7760\n",
            "Epoch 46/75\n",
            "25/25 [==============================] - 3s 122ms/step - accuracy: 0.3251 - loss: 1.8158 - val_accuracy: 0.3368 - val_loss: 1.7702\n",
            "Epoch 47/75\n",
            "25/25 [==============================] - 3s 124ms/step - accuracy: 0.3250 - loss: 1.8125 - val_accuracy: 0.3380 - val_loss: 1.7671\n",
            "Epoch 48/75\n",
            "25/25 [==============================] - 3s 123ms/step - accuracy: 0.3246 - loss: 1.8068 - val_accuracy: 0.3398 - val_loss: 1.7614\n",
            "Epoch 49/75\n",
            "25/25 [==============================] - 3s 126ms/step - accuracy: 0.3250 - loss: 1.8035 - val_accuracy: 0.3381 - val_loss: 1.7568\n",
            "Epoch 50/75\n",
            "25/25 [==============================] - 3s 125ms/step - accuracy: 0.3301 - loss: 1.7974 - val_accuracy: 0.3434 - val_loss: 1.7516\n",
            "Epoch 51/75\n",
            "25/25 [==============================] - 3s 123ms/step - accuracy: 0.3298 - loss: 1.7942 - val_accuracy: 0.3434 - val_loss: 1.7475\n",
            "Epoch 52/75\n",
            "25/25 [==============================] - 3s 124ms/step - accuracy: 0.3316 - loss: 1.7890 - val_accuracy: 0.3410 - val_loss: 1.7465\n",
            "Epoch 53/75\n",
            "25/25 [==============================] - 3s 123ms/step - accuracy: 0.3337 - loss: 1.7863 - val_accuracy: 0.3475 - val_loss: 1.7353\n",
            "Epoch 54/75\n",
            "25/25 [==============================] - 3s 128ms/step - accuracy: 0.3378 - loss: 1.7792 - val_accuracy: 0.3463 - val_loss: 1.7364\n",
            "Epoch 55/75\n",
            "25/25 [==============================] - 3s 124ms/step - accuracy: 0.3395 - loss: 1.7727 - val_accuracy: 0.3520 - val_loss: 1.7285\n",
            "Epoch 56/75\n",
            "25/25 [==============================] - 3s 125ms/step - accuracy: 0.3355 - loss: 1.7734 - val_accuracy: 0.3500 - val_loss: 1.7270\n",
            "Epoch 57/75\n",
            "25/25 [==============================] - 3s 123ms/step - accuracy: 0.3400 - loss: 1.7631 - val_accuracy: 0.3533 - val_loss: 1.7218\n",
            "Epoch 58/75\n",
            "25/25 [==============================] - 3s 124ms/step - accuracy: 0.3446 - loss: 1.7612 - val_accuracy: 0.3539 - val_loss: 1.7202\n",
            "Epoch 59/75\n",
            "25/25 [==============================] - 3s 124ms/step - accuracy: 0.3455 - loss: 1.7578 - val_accuracy: 0.3583 - val_loss: 1.7125\n",
            "Epoch 60/75\n",
            "25/25 [==============================] - 3s 124ms/step - accuracy: 0.3473 - loss: 1.7566 - val_accuracy: 0.3613 - val_loss: 1.7086\n",
            "Epoch 61/75\n",
            "25/25 [==============================] - 3s 123ms/step - accuracy: 0.3486 - loss: 1.7512 - val_accuracy: 0.3641 - val_loss: 1.7023\n",
            "Epoch 62/75\n",
            "25/25 [==============================] - 3s 125ms/step - accuracy: 0.3506 - loss: 1.7446 - val_accuracy: 0.3645 - val_loss: 1.7003\n",
            "Epoch 63/75\n",
            "25/25 [==============================] - 3s 124ms/step - accuracy: 0.3482 - loss: 1.7449 - val_accuracy: 0.3694 - val_loss: 1.6949\n",
            "Epoch 64/75\n",
            "25/25 [==============================] - 3s 125ms/step - accuracy: 0.3538 - loss: 1.7401 - val_accuracy: 0.3691 - val_loss: 1.6922\n",
            "Epoch 65/75\n",
            "25/25 [==============================] - 3s 124ms/step - accuracy: 0.3526 - loss: 1.7373 - val_accuracy: 0.3729 - val_loss: 1.6871\n",
            "Epoch 66/75\n",
            "25/25 [==============================] - 3s 122ms/step - accuracy: 0.3550 - loss: 1.7334 - val_accuracy: 0.3739 - val_loss: 1.6835\n",
            "Epoch 67/75\n",
            "25/25 [==============================] - 3s 125ms/step - accuracy: 0.3591 - loss: 1.7278 - val_accuracy: 0.3774 - val_loss: 1.6783\n",
            "Epoch 68/75\n",
            "25/25 [==============================] - 3s 125ms/step - accuracy: 0.3583 - loss: 1.7254 - val_accuracy: 0.3803 - val_loss: 1.6742\n",
            "Epoch 69/75\n",
            "25/25 [==============================] - 3s 122ms/step - accuracy: 0.3613 - loss: 1.7209 - val_accuracy: 0.3820 - val_loss: 1.6712\n",
            "Epoch 70/75\n",
            "25/25 [==============================] - 3s 123ms/step - accuracy: 0.3603 - loss: 1.7196 - val_accuracy: 0.3862 - val_loss: 1.6627\n",
            "Epoch 71/75\n",
            "25/25 [==============================] - 4s 166ms/step - accuracy: 0.3619 - loss: 1.7178 - val_accuracy: 0.3881 - val_loss: 1.6617\n",
            "Epoch 72/75\n",
            "25/25 [==============================] - 3s 127ms/step - accuracy: 0.3652 - loss: 1.7148 - val_accuracy: 0.3883 - val_loss: 1.6568\n",
            "Epoch 73/75\n",
            "25/25 [==============================] - 3s 129ms/step - accuracy: 0.3641 - loss: 1.7116 - val_accuracy: 0.3906 - val_loss: 1.6544\n",
            "Epoch 74/75\n",
            "25/25 [==============================] - 3s 126ms/step - accuracy: 0.3682 - loss: 1.7092 - val_accuracy: 0.3915 - val_loss: 1.6506\n",
            "Epoch 75/75\n",
            "25/25 [==============================] - 3s 127ms/step - accuracy: 0.3697 - loss: 1.7048 - val_accuracy: 0.3909 - val_loss: 1.6490\n",
            "Network takes 294.997 seconds to train\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}